{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 量子变分自编码器\n",
    "\n",
    "<em> Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved. </em>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 概览\n",
    "\n",
    "在这个案例中，我们将展示如何训练量子自编码器来压缩和重建给定的量子态（混合态）[1]。\n",
    "\n",
    "### 原理\n",
    "\n",
    "量子自编码器的形式与经典自编码器非常相似，同样由编码器 $E$（encoder）和解码器 $D$（decoder）组成。对于输入的 $N$ 量子比特系统的量子态 $\\rho_{in}$（这里我们采用量子力学的密度算符表示来描述混合态），先用编码器 $E = U(\\theta)$ 将信息编码到其中的部分量子比特上。这部分量子比特记为**系统 $A$**。对剩余的量子比特 （这部分记为**系统 $B$**）进行测量并丢弃后，我们就得到了压缩后的量子态 $\\rho_{encode}$！压缩后的量子态维度和量子系统 $A$ 的维度大小保持一致。假设我们需要 $N_A$ 个量子比特来描述系统 $A$ ，那么编码后量子态 $\\rho_{encode}$ 的维度就是 $2^{N_A}\\times 2^{N_A}$。 这里比较特殊的是, 对应我们这一步测量并丢弃的操作的数学操作是 partial trace。读者在直观上可以理解为张量积 $\\otimes$ 的逆运算。\n",
    "\n",
    "我们不妨看一个具体的例子来理解。给定 $N_A$ 个量子比特上的一个量子态 $\\rho_A$ 和 $N_B$ 个量子比特上的一个量子态 $\\rho_B$, 那么由 $A、B$ 两个子系统构成的的整体量子系统 $N = N_A+ N_B$ 的量子态就可以描述为: $\\rho_{AB} = \\rho_A \\otimes \\rho_B$。现在我们让整个量子系统在酉矩阵 $U$ 的作用下演化一段时间后，得到了一个新的量子态 $\\tilde{\\rho_{AB}} = U\\rho_{AB}U^\\dagger$。 那么如果这时候我们只想得到量子子系统 A 所处于的新的量子态 $\\tilde{\\rho_A}$， 我们该怎么做呢？很简单，只需要测量量子子系统 $B$ 然后将其丢弃。运算上这一步由 partial trace 来完成 $\\tilde{\\rho_A} = \\text{Tr}_B (\\tilde{\\rho_{AB}})$。在量桨上，我们可以用内置的 `partial_trace(rho_AB, 2**N_A, 2**N_B, 2)` 指令来完成这一运算。**注意：** 其中最后一个参数为 2，表示我们想丢弃量子系统 $B$ 的量子态。\n",
    "\n",
    "![QA-fig-encoder_pipeline](./figures/QA-fig-encoder_pipeline.png \"**图 1.** 量子变分自编码器流程图\")\n",
    "\n",
    "在讨论完编码的过程后，我们来看看如何解码。为了解码量子态 $\\rho_{encode}$，我们需要引入与系统 $B$ 维度相同的系统 $C$ 并且初始态取为 $|0\\dots0\\rangle$ 态。再用解码器 $D = U^\\dagger(\\theta)$ 作用在整个量子系统 $A+C$ 上对系统 A 中压缩的信息进行解码。我们希望最后得到的量子态 $\\rho_{out}$ 与 $\\rho_{in}$ 尽可能相似并用 Uhlmann-Josza 保真度 $F$ （Fidelity）来衡量他们之间的相似度。\n",
    "\n",
    "$$\n",
    "F(\\rho_{in}, \\rho_{out}) = \\left(\\operatorname{tr} \\sqrt{\\sqrt{\\rho_{in}} \\rho_{out} \\sqrt{\\rho_{in}}} \\right)^{2}.\n",
    "\\tag{1}\n",
    "$$\n",
    "\n",
    "最后，通过优化编码器里的参数，我们就可以尽可能地提高 $\\rho_{in}$ 与 $\\rho_{out}$ 的保真度。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Paddle Quantum 实现\n",
    "\n",
    "下面我们通过一个简单的例子来展示量子自编码器的工作流程和原理。这里我们先引入必要的 package。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:58:04.575047Z",
     "start_time": "2021-04-30T08:58:04.555128Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { white-space: pre !important; }</style>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.core.display import HTML\n",
    "display(HTML(\"<style>pre { white-space: pre !important; }</style>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:58:08.278605Z",
     "start_time": "2021-04-30T08:58:05.042247Z"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numpy import diag\n",
    "import scipy\n",
    "import scipy.stats\n",
    "import paddle\n",
    "from paddle import matmul, trace, kron, real\n",
    "from paddle_quantum.circuit import UAnsatz\n",
    "from paddle_quantum.utils import dagger, state_fidelity, partial_trace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 生成初始态\n",
    "\n",
    "我们考虑 $N = 3$ 个量子比特上的量子态 $\\rho_{in}$。我们先通过编码器将其信息编码到下方的两个量子比特（系统 $A$），之后对第一个量子比特（系统 $B$）测量并丢弃，之后引入一个处于 $|0\\rangle$ 态的量子比特（新的参考系统 $C$）来替代丢弃的量子比特 $B$ 的维度。最后通过解码器，将 A 中压缩的信息复原成 $\\rho_{out}$。在这里，我们假设初态是一个混合态并且 $\\rho_{in}$ 的本征谱为 $\\lambda_i \\in \\{0.4， 0.2， 0.2， 0.1， 0.1, \\,0, \\,0, \\,0 \\}$，然后通过作用一个随机的酉变换来生成初态 $\\rho_{in}$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:58:08.292440Z",
     "start_time": "2021-04-30T08:58:08.281478Z"
    }
   },
   "outputs": [],
   "source": [
    "N_A = 2        # 系统 A 的量子比特数\n",
    "N_B = 1        # 系统 B 的量子比特数\n",
    "N = N_A + N_B  # 总的量子比特数\n",
    "\n",
    "scipy.random.seed(1)                            # 固定随机种子\n",
    "V = scipy.stats.unitary_group.rvs(2**N)         # 随机生成一个酉矩阵\n",
    "D = diag([0.4, 0.2, 0.2, 0.1, 0.1, 0, 0, 0])    # 输入目标态rho的谱\n",
    "V_H = V.conj().T                                # 进行厄尔米特转置\n",
    "rho_in = (V @ D @ V_H).astype('complex128')     # 生成 rho_in\n",
    "\n",
    "# 初始化量子系统 C\n",
    "rho_C = np.diag([1,0]).astype('complex128')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 搭建量子神经网络\n",
    "\n",
    "在这里，我们用量子神经网络来作为编码器和解码器。假设系统 A 有 $N_A$ 个量子比特，系统 $B$ 和 $C$ 各有$N_B$ 个量子比特，量子神经网络的深度为 $D$。编码器 $E$ 作用在系统 A 和 B 共同构成的总系统上，解码器 $D$ 作用在 $A$ 和 $C$ 共同构成的总系统上。在我们的问题里，$N_{A} = 2$，$N_{B} = 1$。\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:58:08.310082Z",
     "start_time": "2021-04-30T08:58:08.303101Z"
    }
   },
   "outputs": [],
   "source": [
    "# 设置电路参数\n",
    "cir_depth = 6                        # 电路深度\n",
    "block_len = 2                        # 每个模组的长度\n",
    "theta_size = N*block_len*cir_depth   # 网络参数 theta 的大小\n",
    "\n",
    "# 搭建编码器 Encoder E\n",
    "def Encoder(theta):\n",
    "\n",
    "    # 用 UAnsatz 初始化网络\n",
    "    cir = UAnsatz(N)\n",
    "    \n",
    "    # 搭建层级结构：\n",
    "    for layer_num in range(cir_depth):\n",
    "        \n",
    "        for which_qubit in range(N):\n",
    "            cir.ry(theta[block_len*layer_num*N + which_qubit], which_qubit)\n",
    "            cir.rz(theta[(block_len*layer_num + 1)*N + which_qubit], which_qubit)\n",
    "\n",
    "        for which_qubit in range(N-1):\n",
    "            cir.cnot([which_qubit, which_qubit + 1])\n",
    "        cir.cnot([N-1, 0])\n",
    "\n",
    "    return cir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 配置训练模型——损失函数\n",
    "\n",
    "在这里，我们定义的损失函数为 \n",
    "\n",
    "$$\n",
    "Loss = 1 - \\langle 0...0|\\rho_{trash}|0...0\\rangle,\n",
    "\\tag{2}\n",
    "$$\n",
    "\n",
    "其中 $\\rho_{trash}$ 是经过编码后丢弃的 $B$ 系统的量子态。接着我们通过飞桨训练量子神经网络，最小化损失函数。如果损失函数最后达到 0，则输入态和输出态完全相同。这就意味着我们完美地实现了压缩和解压缩，这时初态和末态的保真度为  $F(\\rho_{in}, \\rho_{out}) = 1$。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2021-04-30T08:58:31.111139Z",
     "start_time": "2021-04-30T08:58:20.555305Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter: 10 loss: 0.1683 fid: 0.8211\n",
      "iter: 20 loss: 0.1231 fid: 0.8720\n",
      "iter: 30 loss: 0.1122 fid: 0.8810\n",
      "iter: 40 loss: 0.1058 fid: 0.8864\n",
      "iter: 50 loss: 0.1025 fid: 0.8901\n",
      "iter: 60 loss: 0.1019 fid: 0.8907\n",
      "iter: 70 loss: 0.1013 fid: 0.8914\n",
      "iter: 80 loss: 0.1012 fid: 0.8917\n",
      "iter: 90 loss: 0.1010 fid: 0.8921\n",
      "iter: 100 loss: 0.1008 fid: 0.8924\n",
      "\n",
      "训练后的电路：\n",
      "--Ry(3.935)----Rz(2.876)----*---------X----Ry(2.678)----Rz(6.372)----*---------X----Ry(5.516)----Rz(4.082)----*---------X----Ry(1.199)----Rz(1.584)----*---------X----Ry(4.512)----Rz(0.847)----*---------X----Ry(5.038)----Rz(0.564)----*---------X--\n",
      "                            |         |                              |         |                              |         |                              |         |                              |         |                              |         |  \n",
      "--Ry(2.045)----Rz(4.282)----X----*----|----Ry(6.116)----Rz(6.203)----X----*----|----Ry(5.135)----Rz(4.828)----X----*----|----Ry(3.532)----Rz(3.827)----X----*----|----Ry(0.497)----Rz(1.693)----X----*----|----Ry(5.243)----Rz(5.329)----X----*----|--\n",
      "                                 |    |                                   |    |                                   |    |                                   |    |                                   |    |                                   |    |  \n",
      "--Ry(2.706)----Rz(4.168)---------X----*----Ry(2.141)----Rz(2.014)---------X----*----Ry(5.364)----Rz(-0.34)---------X----*----Ry(4.014)----Rz(2.668)---------X----*----Ry(3.419)----Rz(1.952)---------X----*----Ry(4.255)----Rz(1.856)---------X----*--\n",
      "                                                                                                                                                                                                                                                      \n"
     ]
    }
   ],
   "source": [
    "# 超参数设置\n",
    "N_A = 2        # 系统 A 的量子比特数\n",
    "N_B = 1        # 系统 B 的量子比特数\n",
    "N = N_A + N_B  # 总的量子比特数\n",
    "LR = 0.2       # 设置学习速率\n",
    "ITR = 100      # 设置迭代次数\n",
    "SEED = 15      # 固定初始化参数用的随机数种子\n",
    "\n",
    "class NET(paddle.nn.Layer):\n",
    "\n",
    "    def __init__(self, shape, dtype='float64'):\n",
    "        super(NET, self).__init__()\n",
    "        \n",
    "        # 将 Numpy array 转换成 Paddle 中支持的 Tensor\n",
    "        self.rho_in = paddle.to_tensor(rho_in)\n",
    "        self.rho_C = paddle.to_tensor(rho_C)\n",
    "        self.theta = self.create_parameter(shape=shape,\n",
    "                                           default_initializer=paddle.nn.initializer.Uniform(low=0.0, high=2 * np.pi),\n",
    "                                           dtype=dtype, is_bias=False)\n",
    "    \n",
    "    # 定义损失函数和前向传播机制\n",
    "    def forward(self):\n",
    "      \n",
    "        # 生成初始的编码器 E 和解码器 D\n",
    "        cir = Encoder(self.theta)\n",
    "        E = cir.U\n",
    "        E_dagger = dagger(E)\n",
    "        D = E_dagger\n",
    "        D_dagger = E\n",
    "\n",
    "        # 编码量子态 rho_in\n",
    "        rho_BA = matmul(matmul(E, self.rho_in), E_dagger)\n",
    "        \n",
    "        # 取 partial_trace() 获得 rho_encode 与 rho_trash\n",
    "        rho_encode = partial_trace(rho_BA, 2 ** N_B, 2 ** N_A, 1)\n",
    "        rho_trash = partial_trace(rho_BA, 2 ** N_B, 2 ** N_A, 2)\n",
    "\n",
    "        # 解码得到量子态 rho_out\n",
    "        rho_CA = kron(self.rho_C, rho_encode)\n",
    "        rho_out = matmul(matmul(D, rho_CA), D_dagger)\n",
    "        \n",
    "        # 通过 rho_trash 计算损失函数\n",
    "        zero_Hamiltonian = paddle.to_tensor(np.diag([1,0]).astype('complex128'))\n",
    "        loss = 1 - real(trace(matmul(zero_Hamiltonian, rho_trash)))\n",
    "\n",
    "        return loss, self.rho_in, rho_out, cir\n",
    "\n",
    "\n",
    "paddle.seed(SEED)\n",
    "# 生成网络\n",
    "net = NET([theta_size])\n",
    "# 一般来说，我们利用 Adam 优化器来获得相对好的收敛\n",
    "# 当然你可以改成 SGD 或者是 RMS prop.\n",
    "opt = paddle.optimizer.Adam(learning_rate=LR, parameters=net.parameters())\n",
    "\n",
    "# 优化循环\n",
    "for itr in range(1, ITR + 1):\n",
    "    #  前向传播计算损失函数\n",
    "    loss, rho_in, rho_out, cir = net()\n",
    "    # 反向传播极小化损失函数\n",
    "    loss.backward()\n",
    "    opt.minimize(loss)\n",
    "    opt.clear_grad()\n",
    "    # 计算并打印保真度\n",
    "    fid = state_fidelity(rho_in.numpy(), rho_out.numpy())\n",
    "    if itr % 10 == 0:\n",
    "        print('iter:', itr, 'loss:', '%.4f' % loss, 'fid:', '%.4f' % np.square(fid))\n",
    "    if itr == ITR:\n",
    "        print(\"\\n训练后的电路：\") \n",
    "        print(cir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "如果系统 A 的维度为 $d_A$，容易证明量子自编码器能实现的压缩重建最大保真度为 $\\rho_{in}$ 最大的 $d_A$ 个本征值之和。在我们的示例里 $d_A = 4$，理论最大保真度为 \n",
    "\n",
    "$$\n",
    "F_{\\text{max}}(\\rho_{in}, \\rho_{out})  = \\sum_{j=1}^{d_A} \\lambda_j(\\rho_{in})= 0.4 + 0.2 + 0.2 + 0.1 = 0.9.\n",
    "\\tag{3}\n",
    "$$\n",
    "\n",
    "通过 100 次迭代，我们训练出的量子自编码器保真度达到 0.89 以上，和理论最大值非常接近。\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_______\n",
    "\n",
    "## 参考文献\n",
    "\n",
    "[1] Romero, J., Olson, J. P. & Aspuru-Guzik, A. Quantum autoencoders for efficient compression of quantum data. [Quantum Sci. Technol. 2, 045001 (2017).](https://iopscience.iop.org/article/10.1088/2058-9565/aa8072)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
